import numpy as np
from matplotlib import pyplot as plt

import data_read


# 读取数据
def read_data():
    [dataset, label] = data_read.read_in("..\\train_data")
    path_index = np.where(label == 'train_path')[0]
    tree_index = np.where(label == 'train_tree')[0]
    path_feature = dataset[path_index[0]:path_index[-1] + 1]
    tree_feature = dataset[tree_index[0]:tree_index[-1] + 1]
    data_x = np.concatenate((path_feature, tree_feature))[:, [0, 10]]
    data_y = np.append(np.full(np.shape(path_feature)[0], "path"), np.full(np.shape(tree_feature)[0], "tree"))
    return data_x, data_y


# 数据可视化
def vis(x, y):
    plt.figure(figsize=(10, 8))
    tree = np.where(y == "tree")
    path = np.where(y == "path")
    trees, = plt.plot(np.ravel(x[tree, 0]), np.ravel(x[tree, 1]), 'or', markersize=4)
    paths, = plt.plot(np.ravel(x[path, 0]), np.ravel(x[path, 1]), 'og', markersize=4)
    plt.xlabel("X1")
    plt.ylabel("X2")
    plt.legend([trees, paths], ["tree", "path"])
    plt.show()


if __name__ == '__main__':
    data, clazz = read_data()
    vis(data, clazz)

